# @Author：zh
# @Data：2021/12/28 19:53
# @：PyCharm
# Python版本：3.7


## 用于卷积神经网络的 模型预测
## 特别注意： 输入预测的数据也应该进行预变换（像素归一化，分布归一化）

from torch import  nn
from torchvision import transforms
import cv2 as cv
import torch
import  numpy as np

class CNN(nn.Module):
    def __init__(self):
        super(CNN,self).__init__()

        ##  输入数据格式 [N C H W]
        self.layer1 = nn.Sequential(
            ### 参数:输入通道数  输出通道数 卷积核大小
            nn.Conv2d(1,16,kernel_size=3), # 16, [26 ,26]
            nn.BatchNorm2d(16),
            nn.ReLU(inplace=True))

        self.layer2 = nn.Sequential(
            nn.Conv2d(16,32,kernel_size=3),# 32,[ 24, 24]
            nn.BatchNorm2d(32),
            # inplace表示覆盖运算
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2,stride=2)) # 32,[ 12,12 ]    (24-2) /2 +1

        self.layer3 = nn.Sequential(
            nn.Conv2d(32,64,kernel_size=3), # 64,[10,10]
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True))

        self.layer4 = nn.Sequential(
            nn.Conv2d(64,128,kernel_size=3),  # 128,[8,8]
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2,stride=2))  # 128, [4,4]

        self.fc = nn.Sequential(
            nn.Linear(128 * 4 * 4,1024),
            nn.ReLU(inplace=True),
            nn.Linear(1024,128),
            nn.ReLU(inplace=True),
            nn.Linear(128,10))

    def forward(self,x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        ## view与reshape类似 将前面多维度的tensor展平成一维
        ## 参数： 10000=batchSize  128*4*4
        x = x.view(x.size(0),-1)
        x = self.fc(x)
        ## [10000,10]
        return x

# the model class: 你所定义的模型的class
model = CNN()
model.load_state_dict(torch.load("Model_Minst_CNN.pt"))
model.eval()

##初始数据变换
transform_data= transforms.Compose([transforms.ToTensor(),transforms.Normalize([0.5],[0.5])])


img=cv.imread("data/TestData/5.jpg",)
img=cv.cvtColor(img, cv.COLOR_BGR2GRAY)
img=cv.resize(img,(28,28))
# cv.imshow("res",img)
# cv.waitKey(0)
array=np.array(img)

tensor=transform_data(array)   ## 在完成归一化的同时，将array转化为了tensor

# 将输入的图像从array格式转为image

# 因为训练的时候输入的照片的维度为(batch_size, channel, H, W) 所以需要再新增一个维度
# 增加的维度是batch size，这里输入的一张图片，所以为1


tensor = tensor.view(1, 1, 28, 28)
#print(tensor)

output = model(tensor.to(torch.float32))
_, prediction = torch.max(output, 1)
#将预测结果从tensor转为array，并抽取结果
prediction = prediction.numpy()[0]

print(prediction)


###  为了保证准确率： 输入图片应该是 黑底白字



